package cn.ywyself.extend.utils;

import org.apache.commons.codec.binary.Base64;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.HashMap;
import java.util.Map;

/**
 * RSA签名工具
 *
 * @author : ywyself
 * @created : 2018-10-22 9:04
 */
public class RSAUtils {
    public static final String SHA1_WITH_RSA = "SHA1WithRSA";
    public static final String SHA256_WITH_RSA = "SHA256WithRSA";

    private static final String RSA = "RSA";
    
    private static final String CHARSET = "UTF-8";

    private static final int MAX_ENCRYPT_BLOCK = 117;

    public RSAUtils() {
    }

    /**
     * 公钥加密
     *
     * @param data      加密数据
     * @param strPubKey 文本公钥
     * @return 加密后的内容
     */
    public static byte[] encryptByPublicKey(byte[] data, String strPubKey) throws NoSuchAlgorithmException, InvalidKeySpecException, NoSuchPaddingException, InvalidKeyException, BadPaddingException, IllegalBlockSizeException, IOException {
        byte[] keyBytes = Base64.decodeBase64(strPubKey);
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        Key publicK = keyFactory.generatePublic(x509KeySpec);
        // 对数据加密
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.ENCRYPT_MODE, publicK);
        int inputLen = data.length;
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        int offSet = 0;
        byte[] cache;
        int i = 0;
        // 对数据分段加密
        while (inputLen - offSet > 0) {
            if (inputLen - offSet > MAX_ENCRYPT_BLOCK) {
                cache = cipher.doFinal(data, offSet, MAX_ENCRYPT_BLOCK);
            } else {
                cache = cipher.doFinal(data, offSet, inputLen - offSet);
            }
            out.write(cache, 0, cache.length);
            i++;
            offSet = i * MAX_ENCRYPT_BLOCK;
        }
        byte[] encryptedData = out.toByteArray();
        out.close();
        return encryptedData;
    }

    /**
     * 公钥加密
     *
     * @param data      加密数据
     * @param strPubKey 文本公钥
     * @return 加密后的内容
     */
    public static String encryptByPublicKey(String data, String strPubKey) throws IOException, NoSuchPaddingException, NoSuchAlgorithmException, IllegalBlockSizeException, BadPaddingException, InvalidKeyException, InvalidKeySpecException {
        return Base64.encodeBase64String(encryptByPublicKey(data.getBytes(CHARSET), strPubKey));
    }

    /**
     * RSA签名验证
     *
     * @param data               待验证的内容
     * @param sign               签名内容
     * @param publicKey          公钥
     * @param signatureAlgorithm 签名算法
     * @return true|验证通过
     */
    public static boolean verify(byte[] data, byte[] sign, PublicKey publicKey, String signatureAlgorithm) throws NoSuchAlgorithmException, InvalidKeyException, SignatureException {
        Signature signature = Signature.getInstance(signatureAlgorithm);
        signature.initVerify(publicKey);
        signature.update(data);
        return signature.verify(sign);
    }

    /**
     * RSA签名验证
     *
     * @param data               待验证的内容
     * @param sign               签名内容
     * @param strPubKey          文本公钥
     * @param signatureAlgorithm 签名算法
     * @return true|验证通过
     */
    public static boolean verify(String data, String sign, String strPubKey, String signatureAlgorithm) throws InvalidKeySpecException, NoSuchAlgorithmException, SignatureException, InvalidKeyException, UnsupportedEncodingException {
        return verify(data.getBytes(CHARSET), Base64.decodeBase64(sign), getPublicKey(strPubKey), signatureAlgorithm);
    }

    /**
     * RSA签名计算
     *
     * @param data               待签名内容
     * @param privateKey         私钥
     * @param signatureAlgorithm 签名算法
     * @return 签名
     */
    public static byte[] sign(byte[] data, PrivateKey privateKey, String signatureAlgorithm) throws NoSuchAlgorithmException, InvalidKeyException, SignatureException {
        Signature signature = Signature.getInstance(signatureAlgorithm);
        signature.initSign(privateKey);
        signature.update(data);
        return signature.sign();
    }

    /**
     * RSA签名计算
     *
     * @param data               待签名内容
     * @param strPriKey          文本私钥
     * @param signatureAlgorithm 签名算法
     * @return 签名
     */
    public static String sign(String data, String strPriKey, String signatureAlgorithm) throws InvalidKeyException, NoSuchAlgorithmException, SignatureException, InvalidKeySpecException, UnsupportedEncodingException {
        return Base64.encodeBase64String(sign(data.getBytes(CHARSET), getPrivateKey(strPriKey), signatureAlgorithm));
    }

    /**
     * 将 文本公钥 转换为 公钥对象
     *
     * @param strPubKey 文本公钥
     * @return 公钥对象
     */
    public static PublicKey getPublicKey(String strPubKey) throws NoSuchAlgorithmException, InvalidKeySpecException {
        X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(Base64.decodeBase64(strPubKey));
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        return keyFactory.generatePublic(pubKeySpec);
    }

    /**
     * 将 文本私钥 转换为 私钥对象
     *
     * @param strPriKey 文本私钥
     * @return 私钥对象
     */
    public static PrivateKey getPrivateKey(String strPriKey) throws NoSuchAlgorithmException, InvalidKeySpecException {
        PKCS8EncodedKeySpec priKeySpec = new PKCS8EncodedKeySpec(Base64.decodeBase64(strPriKey));
        KeyFactory keyFactory = KeyFactory.getInstance(RSA);
        return keyFactory.generatePrivate(priKeySpec);
    }

    // 以上为方法 ****************************************************************** //

    /**
     * 将公钥对象转换为公钥文本
     *
     * @param pubKey 公钥
     * @return 公钥文本
     */
    public static String getBase64PublicKeyString(PublicKey pubKey) {
        return Base64.encodeBase64String(pubKey.getEncoded()).trim();
    }

    /**
     * 将私钥对象转换为私钥文本
     *
     * @param priKey 私钥
     * @return 私钥文本
     */
    public static String getBase64PrivateKeyString(PrivateKey priKey) {
        return Base64.encodeBase64String(priKey.getEncoded()).trim();
    }

    private static final int RSA_SIZE_2048 = 2048;
    private static final int RSA_SIZE_1024 = 1024;
    private static final String PUBLIC_KEY = "publicKey";
    private static final String PUBLIC_KEY_B64 = "publicKeyBase64";
    private static final String PRIVATE_KEY = "privateKey";
    private static final String PRIVATE_KEY_B64 = "privateKeyBase64";

    /**
     * 创建公钥，私钥
     *
     * @param keySize 长度
     * @return 公钥，私钥，以及他们的Base64编码后的内容
     * @throws NoSuchAlgorithmException 没有这个算法
     */
    private static Map<String, Object> createKey(int keySize) throws NoSuchAlgorithmException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance(RSA);
        keyGen.initialize(keySize, new SecureRandom());
        KeyPair key = keyGen.generateKeyPair();
        PublicKey pubKey = key.getPublic();
        PrivateKey priKey = key.getPrivate();
        Map<String, Object> map = new HashMap<>();
        map.put(PUBLIC_KEY, pubKey);
        map.put(PRIVATE_KEY, priKey);
        map.put(PUBLIC_KEY_B64, getBase64PublicKeyString(pubKey));
        map.put(PRIVATE_KEY_B64, getBase64PrivateKeyString(priKey));
        return map;
    }

    /**
     * 创建公钥，私钥
     *
     * @return 1024 公钥，私钥
     */
    public static Map<String, Object> create1024Key() throws NoSuchAlgorithmException {
        return createKey(RSA_SIZE_1024);
    }

    /**
     * 创建公钥，私钥
     *
     * @return 2048 公钥，私钥
     */
    public static Map<String, Object> create2048Key() throws NoSuchAlgorithmException {
        return createKey(RSA_SIZE_2048);
    }

    public static void main(String[] args) throws Exception {
        Map<String, Object> key = create2048Key();
        PublicKey pubKey = (PublicKey) key.get(PUBLIC_KEY);
        PrivateKey priKey = (PrivateKey) key.get(PRIVATE_KEY);
        String strPubKey = (String) key.get(PUBLIC_KEY_B64);
        String strPriKey = (String) key.get(PRIVATE_KEY_B64);
        System.out.println("公钥：" + strPubKey);
        System.out.println("私钥：" + strPriKey);
        String data = "123456";
        String signData = sign(data, strPriKey, SHA256_WITH_RSA);
        System.out.println("签名值为：" + signData);
        boolean result = verify(data, signData, strPubKey, SHA256_WITH_RSA);
        System.out.println("验签结果为：" + result);
        System.out.println();
    }
}
